package cn.edu.baiyunu.myapplication.utils

import android.content.Context
import android.graphics.Bitmap
import android.graphics.Canvas
import android.graphics.Color
import android.graphics.Matrix
import android.graphics.Paint
import android.graphics.PorterDuff
import android.graphics.PorterDuffXfermode
import android.util.Log
import cn.edu.baiyunu.myapplication.utils.MessageType
import cn.edu.baiyunu.myapplication.utils.showMessage
import ai.onnxruntime.OnnxTensor
import ai.onnxruntime.OrtEnvironment
import ai.onnxruntime.OrtSession
import java.io.File
import java.io.FileOutputStream
import java.io.IOException
import java.nio.ByteBuffer
import java.nio.FloatBuffer

/**
 * MODNet人像抠图处理类
 * 基于MODNet: Real-Time Trimap-Free Portrait Matting via Objective Decomposition
 * https://github.com/ZHKKKe/MODNet
 */
class MODNetProcessor(private val context: Context) {
    private var ortEnvironment: OrtEnvironment? = null
    private var ortSession: OrtSession? = null
    private val TAG = "MODNetProcessor"
    
    // MODNet模型的输入尺寸
    private val INPUT_WIDTH = 512
    private val INPUT_HEIGHT = 512
    
    // 归一化参数
    private val MEAN = floatArrayOf(0.485f, 0.456f, 0.406f)
    private val STD = floatArrayOf(0.229f, 0.224f, 0.225f)
    
    /**
     * 初始化MODNet模型
     */
    fun initialize() {
        try {
            Log.d(TAG, "开始初始化MODNet模型...")
            
            // 从assets复制模型到本地存储
            val modelFile = File(context.filesDir, "modnet.onnx")
            
            if (!modelFile.exists()) {
                Log.d(TAG, "本地缓存中没有找到模型文件，尝试从assets加载")
                
                try {
                    // 我们已经确认模型文件在assets目录下
                    val modelPath = "modnet.onnx"
                    context.assets.open(modelPath).use { input ->
                        Log.d(TAG, "成功打开assets模型文件: $modelPath")
                        FileOutputStream(modelFile).use { output ->
                            val bytes = input.copyTo(output)
                            Log.d(TAG, "已将模型(${bytes}字节)复制到: ${modelFile.absolutePath}")
                        }
                    }
                } catch (e: IOException) {
                    Log.e(TAG, "复制模型文件失败: ${e.message}")
                    e.printStackTrace()
                    showModelMissingMessage()
                    return
                }
            } else {
                Log.d(TAG, "本地已缓存模型文件: ${modelFile.absolutePath}")
            }
            
            // 检查模型文件是否存在且大小正确
            if (!modelFile.exists() || modelFile.length() == 0L) {
                Log.e(TAG, "模型文件不存在或大小为0: ${modelFile.absolutePath}")
                showModelMissingMessage()
                return
            }
            
            Log.d(TAG, "模型文件准备完成，大小: ${modelFile.length()}字节")
            
            // 初始化ONNX Runtime环境
            try {
                ortEnvironment = OrtEnvironment.getEnvironment()
                Log.d(TAG, "ONNX Runtime环境初始化成功")
            } catch (e: UnsatisfiedLinkError) {
                Log.e(TAG, "ONNX Runtime库加载失败，请检查依赖是否正确添加: ${e.message}")
                e.printStackTrace()
                showModelMissingMessage()
                return
            } catch (e: NoClassDefFoundError) {
                Log.e(TAG, "找不到ONNX Runtime类定义，请检查依赖: ${e.message}")
                e.printStackTrace()
                showModelMissingMessage()
                return
            }
            
            // 创建ONNX会话选项
            val sessionOptions = OrtSession.SessionOptions()
            sessionOptions.setOptimizationLevel(OrtSession.SessionOptions.OptLevel.ALL_OPT)
            
            // 创建ONNX会话
            ortSession = ortEnvironment?.createSession(modelFile.absolutePath, sessionOptions)
            
            if (ortSession != null) {
                val inputNames = ortSession?.inputNames?.joinToString() ?: "未知"
                val outputNames = ortSession?.outputNames?.joinToString() ?: "未知"
                Log.d(TAG, "MODNet ONNX模型加载成功! 输入节点: $inputNames, 输出节点: $outputNames")
            } else {
                Log.e(TAG, "ONNX会话创建失败，session为null")
                showModelMissingMessage()
            }
        } catch (e: Exception) {
            Log.e(TAG, "MODNet ONNX模型加载失败: ${e.message}")
            e.printStackTrace()
            showModelMissingMessage()
        }
    }
    
    /**
     * 显示模型缺失提示
     */
    private fun showModelMissingMessage() {
        showMessage(
            "MODNet模型文件缺失，请将模型文件放入assets目录",
            MessageType.ERROR
        )
    }
    
    /**
     * 处理图像，生成人像蒙版
     * @param bitmap 输入的原始图像
     * @return 人像蒙版 (Alpha通道)
     */
    fun processImage(bitmap: Bitmap): Bitmap? {
        if (ortSession == null || ortEnvironment == null) {
            Log.e(TAG, "MODNet模型未加载，无法处理图像")
            return null
        }
        
        var inputTensor: OnnxTensor? = null
        var output: OrtSession.Result? = null
        
        try {
            Log.d(TAG, "开始处理图像，原始尺寸: ${bitmap.width}x${bitmap.height}")
            
            // 转换为ONNX输入格式
            val floatBuffer = bitmapToFloatBuffer(bitmap) // 此方法内部已处理缩放
            
            // 创建输入张量 (1,3,H,W) NCHW格式
            val shape = longArrayOf(1, 3, INPUT_HEIGHT.toLong(), INPUT_WIDTH.toLong())
            inputTensor = OnnxTensor.createTensor(ortEnvironment, floatBuffer, shape)
            
            // 获取模型输入节点名称
            val inputNames = ortSession?.inputNames
            if (inputNames == null || !inputNames.iterator().hasNext()) {
                Log.e(TAG, "获取模型输入节点名称失败")
                return null
            }
            
            val inputName = inputNames.iterator().next()
            Log.d(TAG, "使用输入节点: $inputName 运行模型推理")
            
            // 创建输入映射并执行推理
            val inputs = HashMap<String, OnnxTensor>()
            inputs[inputName] = inputTensor
            output = ortSession?.run(inputs)
            
            if (output == null) {
                Log.e(TAG, "模型推理结果为空")
                return null
            }
            
            // 尝试获取输出张量
            val outputTensor = try {
                // 尝试通过索引获取
                output.get(0) as? OnnxTensor
            } catch (e: Exception) {
                Log.e(TAG, "通过索引获取输出张量失败: ${e.message}")
                try {
                    // 尝试通过名称获取
                    val names = ortSession?.outputNames
                    if (names != null && names.iterator().hasNext()) {
                        val outputName = names.iterator().next()
                        Log.d(TAG, "尝试通过名称获取输出节点: $outputName")
                        output.get(outputName) as? OnnxTensor
                    } else {
                        null
                    }
                } catch (e2: Exception) {
                    Log.e(TAG, "通过名称获取输出张量也失败: ${e2.message}")
                    null
                }
            }
            
            if (outputTensor == null) {
                Log.e(TAG, "输出张量为空或类型不匹配")
                return null
            }
            
            Log.d(TAG, "模型推理完成")
            
            // 转为Bitmap
            val outputMask = onnxTensorToMask(outputTensor, INPUT_WIDTH, INPUT_HEIGHT)
            
            // 将蒙版调整为原始图像尺寸
            val resultMask = resizeBitmap(outputMask, bitmap.width, bitmap.height)
            Log.d(TAG, "成功生成人像蒙版")
            
            return resultMask
            
        } catch (e: Exception) {
            Log.e(TAG, "处理图像失败: ${e.message}")
            e.printStackTrace()
            return null
        } finally {
            // 释放资源
            try {
                inputTensor?.close()
                output?.close()
            } catch (e: Exception) {
                Log.e(TAG, "关闭资源时出错: ${e.message}")
            }
        }
    }
    
    /**
     * 将Bitmap转换为FloatBuffer，适用于ONNX Runtime
     * 预处理图像：缩放、归一化，转换为NCHW格式
     */
    private fun bitmapToFloatBuffer(bitmap: Bitmap): FloatBuffer {
        try {
            Log.d(TAG, "开始预处理图像：${bitmap.width}x${bitmap.height} -> ${INPUT_WIDTH}x${INPUT_HEIGHT}")
            
            // 确保输入图像尺寸正确
            val scaledBitmap = if (bitmap.width != INPUT_WIDTH || bitmap.height != INPUT_HEIGHT) {
                Bitmap.createScaledBitmap(bitmap, INPUT_WIDTH, INPUT_HEIGHT, true)
            } else {
                bitmap
            }
            
            // 为3通道(RGB)分配空间: (1,3,H,W)
            val channelSize = INPUT_HEIGHT * INPUT_WIDTH
            val buffer = FloatBuffer.allocate(3 * channelSize)
            
            // 获取像素数据
            val pixels = IntArray(channelSize)
            scaledBitmap.getPixels(pixels, 0, INPUT_WIDTH, 0, 0, INPUT_WIDTH, INPUT_HEIGHT)
            
            // 对三个通道分别进行归一化处理 (NCHW格式)
            var idx = 0
            
            // 对三个通道(R,G,B)分别处理
            for (c in 0 until 3) {
                val channelOffset = c * channelSize
                
                for (y in 0 until INPUT_HEIGHT) {
                    for (x in 0 until INPUT_WIDTH) {
                        idx = y * INPUT_WIDTH + x
                        
                        // 获取RGB值 (0-255)
                        val pixelValue = when (c) {
                            0 -> (pixels[idx] shr 16) and 0xFF // R
                            1 -> (pixels[idx] shr 8) and 0xFF  // G
                            2 -> pixels[idx] and 0xFF          // B
                            else -> 0
                        }
                        
                        // 归一化: (pixel/255 - mean) / std
                        val normalizedValue = (pixelValue / 255.0f - MEAN[c]) / STD[c]
                        
                        // 在NCHW格式中的索引
                        buffer.put(channelOffset + idx, normalizedValue)
                    }
                }
            }
            
            // 确保position在开始位置
            buffer.rewind()
            
            if (scaledBitmap != bitmap) {
                scaledBitmap.recycle() // 回收临时创建的Bitmap
            }
            
            Log.d(TAG, "图像预处理完成")
            return buffer
            
        } catch (e: Exception) {
            Log.e(TAG, "图像预处理失败: ${e.message}")
            e.printStackTrace()
            throw e
        }
    }
    
    /**
     * 将ONNX输出张量转换为蒙版图像
     */
    private fun onnxTensorToMask(tensor: Any?, width: Int, height: Int): Bitmap {
        val maskBitmap = Bitmap.createBitmap(width, height, Bitmap.Config.ARGB_8888)
        
        // 处理输出张量数据
        if (tensor is OnnxTensor) {
            try {
                // 获取输出数据
                val outputShape = tensor.info.shape
                val floatBuffer = tensor.floatBuffer
                
                // 将输出数据填充到蒙版中
                // MODNet输出通常为[1, 1, height, width]格式
                for (y in 0 until height) {
                    for (x in 0 until width) {
                        // 计算在缓冲区中的索引，通常是y*width + x对于单通道输出
                        var index = y * width + x
                        
                        // 如果模型输出包含批次和通道维度
                        if (outputShape.size >= 4) {
                            // 1(batch) * 1(channel) * height * width + y * width + x
                            index = y * width + x
                        }
                        
                        // 确保索引在范围内
                        if (index < floatBuffer.capacity()) {
                            val alphaValue = floatBuffer.get(index)
                            
                            // 将值从0-1范围转换为0-255
                            val alpha = (alphaValue * 255).toInt().coerceIn(0, 255)
                            
                            // 设置像素
                            maskBitmap.setPixel(x, y, Color.argb(alpha, 255, 255, 255))
                        }
                    }
                }
                
                Log.d(TAG, "成功生成蒙版，尺寸: ${width}x${height}")
            } catch (e: Exception) {
                Log.e(TAG, "处理输出张量失败: ${e.message}")
                e.printStackTrace()
            }
        } else {
            Log.e(TAG, "输出张量格式不正确: ${tensor?.javaClass?.name}")
        }
        
        return maskBitmap
    }
    
    /**
     * 应用蒙版提取人像，生成带透明背景的图像
     * @param bitmap 原始图像
     * @param mask 人像蒙版
     * @return 带透明背景的人像
     */
    fun applyMask(bitmap: Bitmap, mask: Bitmap): Bitmap {
        try {
            Log.d(TAG, "开始应用蒙版提取人像")
            
            // 确保原始图像和蒙版尺寸一致
            if (bitmap.width != mask.width || bitmap.height != mask.height) {
                Log.w(TAG, "原始图像和蒙版尺寸不一致，调整蒙版尺寸")
                val resizedMask = resizeBitmap(mask, bitmap.width, bitmap.height)
                mask.recycle() // 回收不需要的mask
                return applyMask(bitmap, resizedMask) // 递归调用自身
            }
            
            // 创建带Alpha通道的输出Bitmap
            val outputBitmap = Bitmap.createBitmap(bitmap.width, bitmap.height, Bitmap.Config.ARGB_8888)
            val canvas = Canvas(outputBitmap)
            
            // 使用清晰的绘制模式
            val paint = Paint()
            paint.isAntiAlias = true
            paint.isFilterBitmap = true
            
            // 清空画布并设置为完全透明
            canvas.drawColor(Color.TRANSPARENT, PorterDuff.Mode.CLEAR)
            
            // 设置PorterDuff模式为SRC_IN - 只在蒙版区域内绘制源图像
            paint.xfermode = PorterDuffXfermode(PorterDuff.Mode.SRC_IN)
            
            // 先绘制蒙版
            canvas.drawBitmap(mask, 0f, 0f, null)
            
            // 再绘制源图像
            canvas.drawBitmap(bitmap, 0f, 0f, paint)
            
            Log.d(TAG, "成功提取人像")
            return outputBitmap
        } catch (e: Exception) {
            Log.e(TAG, "应用蒙版失败: ${e.message}")
            e.printStackTrace()
            
            // 失败时返回原图
            return bitmap.copy(bitmap.config ?: Bitmap.Config.ARGB_8888, true)
        }
    }
    
    /**
     * 检查MODNet是否已初始化
     */
    fun isInitialized(): Boolean {
        return ortSession != null
    }
    
    /**
     * 更换背景颜色
     * @param bitmap 输入位图（需要带有透明背景）
     * @param backgroundColor 背景颜色
     * @return 更换背景后的位图
     */
    fun changeBackground(bitmap: Bitmap, backgroundColor: Int): Bitmap? {
        if (ortSession == null) {
            Log.e(TAG, "changeBackground: MODNet未初始化")
            return null
        }
        
        try {
            // 创建一个新的位图，防止修改原图
            val width = bitmap.width
            val height = bitmap.height
            
            // 使用ARGB_8888确保支持透明度
            val result = Bitmap.createBitmap(width, height, Bitmap.Config.ARGB_8888)
            
            // 填充背景色
            val canvas = Canvas(result)
            canvas.drawColor(backgroundColor)
            
            // 在背景上绘制原图（带透明度）
            val paint = Paint().apply {
                isAntiAlias = true
                isDither = true
                isFilterBitmap = true
            }
            
            // 绘制原图到新画布上
            canvas.drawBitmap(bitmap, 0f, 0f, paint)
            
            Log.d(TAG, "背景颜色更换成功: ${String.format("#%06X", 0xFFFFFF and backgroundColor)}")
            
            return result
        } catch (e: Exception) {
            Log.e(TAG, "更换背景颜色时出错: ${e.message}")
            e.printStackTrace()
            return null
        }
    }
    
    /**
     * 调整图像尺寸
     */
    private fun resizeBitmap(bitmap: Bitmap, width: Int, height: Int): Bitmap {
        return Bitmap.createScaledBitmap(bitmap, width, height, true)
    }
    
    /**
     * 释放资源
     */
    fun close() {
        try {
            ortSession?.close()
            ortEnvironment?.close()
        } catch (e: Exception) {
            Log.e(TAG, "关闭ONNX资源失败: ${e.message}")
        }
        ortSession = null
        ortEnvironment = null
    }
} 